import os

# os.environ['KMP_DUPLICATE_LIB_OK']='True'
import numpy as np
import pandas as pd
import mplfinance as mpf
import torch
from PIL import Image
from matplotlib import pyplot as plt

# import talib

from dataset_from_csv import DatasetFromCSV
from util.constant_util import BTC_CSV_PATH

mplfinance_image_path = os.path.join("image","mplfinance_image")
os.makedirs(mplfinance_image_path, exist_ok=True)

def protract_k_line(data,savefig_path='stock_price_candlestick.png'):

    df = pd.DataFrame(data,columns=['Date','Open','High','Low','Close','Volume']) # 将ndarray转换为pandas的DataFrame
    df['Date'] = pd.to_datetime(df['Date'],unit='s') # 将时间戳转换为时间  unit='s'表示时间戳是以秒为单位的 # .dt.strftime('%Y-%m-%d %H:%M:%S')
    df = df.set_index('Date')
    # sma_10 = talib.SMA(np.array(df['Close']), 10)
    # sma_30 = talib.SMA(np.array(df['Close']), 30)
    # print(df.head(10))
    # print(df.tail(10))
    # df.index.name = 'Date'
    # 创建一个自定义的面板样式
    addplot = mpf.make_addplot(df['Close'].pct_change(), type='bar', secondary_y=False, color='r', alpha=0.4)

    my_color = mpf.make_marketcolors(up='r',down='g',edge='inherit',wick='inherit',volume='inherit')

    # 绘制K线图并添加自定义功能
    mpf_style = mpf.make_mpf_style(
        # base_mpf_style='charles',
        rc={
            'font.size': 8,
            'ytick.labelsize':10,
            'xtick.labelsize':6,
            'grid.color':'gray',
            'grid.linestyle':'--',
            'grid.linewidth':0.5
        },
        marketcolors=my_color,
        figcolor='(0.82, 0.83, 0.85)',
        gridcolor='(0.82, 0.83, 0.85)',
    )
    # mpf_style.rc['ytick.labelsize'] = 10
    # 设置mplfinance的蜡烛颜色，up为阳线颜色，down为阴线颜色

    figsize = (8, 4)
    # figratio = (3, 2)

    fig, ax_list = mpf.plot(df,
                 type='candle',
                 addplot=addplot,
                 # marketcolors=my_color,
                 mav=(5, 10, 20),
                 style=mpf_style,
                 datetime_format='%d %H:%M',
                 title='Stock Price with Change',
                 ylabel='Price',
                 volume=True,  # 添加成交量图（如果有成交量数据）
                 panel_ratios=(8, 2),  # 设置K线图和成交量图的比例
                 figsize=figsize,  # 设置整个图表的比例
                 tight_layout=True,
                 returnfig=True # 返回图片
                 )

    # 访问y轴并设置次要刻度的可见性
    if ax_list is not None:
        ax = ax_list[0]
        ax.yaxis.set_minor_locator(plt.MaxNLocator(nbins=20))  # 设置次要刻度的数量（可选）
        # ax.yaxis.tick_minor_top(True)  # 如果需要在y轴顶部也显示次要刻度（可选）
        # ax.yaxis.tick_minor_bottom(True)  # 确保y轴底部显示次要刻度（通常默认就是True）
        ax.yaxis.grid(True, which='minor', linestyle='--', linewidth=0.5)
    if savefig_path is not None:
        # 保存图表为文件
        fig.savefig(savefig_path)
    else:
        fig.show()
    # plt.close(fig)  # 显式关闭图形


def merge_pictures(input_image_path_list,ouput_image_path,direction=1):
    images = [Image.open(f) for f in input_image_path_list]
    widths, heights = zip(*(i.size for i in images))

    if direction == 1:
        total_width = sum(widths)
        max_height = max(heights)
    elif direction == 2:
        total_width = max(widths)
        max_height = sum(heights)

    new_im = Image.new('RGB', (total_width, max_height))

    if direction == 1:
        x_offset = 0
        for im in images:
            new_im.paste(im, (x_offset, 0))
            x_offset += im.width
    else:
        y_offset = 0
        for im in images:
            new_im.paste(im, (0, y_offset))
            y_offset += im.height

    # new_im.show()  # 显示拼接后的图像
    new_im.save(ouput_image_path)  # 保存拼接后的图像

def get_forecast_comparison_chart_numpy(known_data,predict_data,reality_data):
    predict_data = predict_data.reshape(predict_data.shape[0],reality_data.shape[1],-1);
    predict_data = np.concatenate((reality_data[:,:,0].reshape(reality_data.shape[0],reality_data.shape[1],1), predict_data), axis=2)
    predict_data = np.concatenate((known_data,predict_data),axis=1)
    reality_data = np.concatenate((known_data,reality_data),axis=1)
    predict_k_line_0_path = os.path.join(mplfinance_image_path,"predict_0_price.png");
    reality_k_line_0_path = os.path.join(mplfinance_image_path,"reality_0_price.png");
    protract_k_line(predict_data[0], savefig_path=predict_k_line_0_path)
    protract_k_line(reality_data[0], savefig_path=reality_k_line_0_path)


    ouput_0_image_path = os.path.join(mplfinance_image_path,"forecast_comparison_chart_0.png")
    merge_pictures([predict_k_line_0_path, reality_k_line_0_path], ouput_0_image_path)

    predict_k_line_last_path = os.path.join(mplfinance_image_path,"predict_last_price.png");
    reality_k_line_last_path = os.path.join(mplfinance_image_path,"reality_last_price.png");
    protract_k_line(predict_data[-1], savefig_path=predict_k_line_last_path)
    protract_k_line(reality_data[-1], savefig_path=reality_k_line_last_path)

    ouput_last_image_path = os.path.join(mplfinance_image_path,"forecast_comparison_chart_last.png")
    merge_pictures([predict_k_line_last_path, reality_k_line_last_path], ouput_last_image_path)

    ouput_image_path = os.path.join(mplfinance_image_path,"forecast_comparison_chart.png")
    merge_pictures([ouput_0_image_path, ouput_last_image_path], ouput_image_path,direction=2)
    # plt.close()
    plt.close('all')
    return np.array(Image.open(ouput_image_path));

def get_forecast_comparison_chart_torch(known_data,predict_data,reality_data):
    torch.set_printoptions(profile="full") # debug 时tensor显示不全和不显示科学计数法的方法
    torch.set_printoptions(sci_mode=False) # debug 时tensor显示不全和不显示科学计数法的方法
    predict_data = predict_data.reshape(predict_data.shape[0],reality_data.shape[1],-1);
    predict_data = torch.cat((reality_data[:,:,0].reshape(reality_data.shape[0],reality_data.shape[1],1),predict_data), dim=2)
    predict_data = torch.cat((known_data, predict_data), dim=1)
    reality_data = torch.cat((known_data, reality_data), dim=1)
    predict_k_line_0_path = os.path.join(mplfinance_image_path,"predict_0_price.png");
    reality_k_line_0_path = os.path.join(mplfinance_image_path,"reality_0_price.png");
    protract_k_line(predict_data[0].detach().cpu().numpy(), savefig_path=predict_k_line_0_path)
    protract_k_line(reality_data[0].detach().cpu().numpy(), savefig_path=reality_k_line_0_path)


    ouput_0_image_path = os.path.join(mplfinance_image_path,"forecast_comparison_chart_0.png")
    merge_pictures([predict_k_line_0_path, reality_k_line_0_path], ouput_0_image_path)

    predict_k_line_last_path = os.path.join(mplfinance_image_path,"predict_last_price.png");
    reality_k_line_last_path = os.path.join(mplfinance_image_path,"reality_last_price.png");
    protract_k_line(predict_data[-1].detach().cpu().numpy(), savefig_path=predict_k_line_last_path)
    protract_k_line(reality_data[-1].detach().cpu().numpy(), savefig_path=reality_k_line_last_path)

    ouput_last_image_path = os.path.join(mplfinance_image_path,"forecast_comparison_chart_last.png")
    merge_pictures([predict_k_line_last_path, reality_k_line_last_path], ouput_last_image_path)

    ouput_image_path = os.path.join(mplfinance_image_path,"forecast_comparison_chart.png")
    merge_pictures([ouput_0_image_path, ouput_last_image_path], ouput_image_path,direction=2)
    # plt.close()
    plt.close('all')
    return np.array(Image.open(ouput_image_path));



if __name__ == '__main__':
    sequence_length = 60 * 24
    train_data = DatasetFromCSV(BTC_CSV_PATH, sequence_length, 30,
                                training_ratio=0.7)
    protract_k_line_1_path = os.path.join(mplfinance_image_path,"stock_price_candlestick_1.png");
    protract_k_line_2_path = os.path.join(mplfinance_image_path,"stock_price_candlestick_2.png");
    protract_k_line(train_data.data.iloc[1000000:1000050].values,savefig_path=protract_k_line_1_path)
    protract_k_line(train_data.data.iloc[1001000:1001050].values,savefig_path=protract_k_line_2_path)

    ouput_image_path = os.path.join(mplfinance_image_path,"ouput_stock_price_candlestick.png")
    merge_pictures([protract_k_line_1_path,protract_k_line_2_path],ouput_image_path)

